#include "llvm/Pass.h"
#include "llvm/IR/Function.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/IntrinsicInst.h"
#include <set>
#include <cxxabi.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/file.h>

using namespace llvm;
using namespace std;

namespace {
    struct lookinPass : public FunctionPass {
        static char ID;
        lookinPass() : FunctionPass(ID) {
        }

	// Create a global string value which is a inbound getelementptr inst.
	// Something like that:
	// @0 = private unnamed_addr constant [4 x i8] c"gcd\00", align 1
	Value* CreateString(StringRef str, LLVMContext &ctx, Module &module) {
		IRBuilder<> builder(ctx);
		// Create a string value.
		auto strVal = builder.CreateGlobalString(str,"",0,&module);
		if (strVal == nullptr) {
			errs() << "Cannot create string: " << str << "\n";
			return nullptr;
		}
		// Wrap the value with getelementptr.
		auto len = cast<PointerType>(strVal->getType())->getElementType()->getArrayNumElements();
		vector<Value*> idxVec;
		idxVec.push_back(CreateInt(0, ctx));
		idxVec.push_back(CreateInt(0, ctx));
		return builder.CreateInBoundsGEP(ArrayType::get(Type::getInt8Ty(ctx),len), strVal,idxVec);
	}

	// Create a global integer value.
	Value* CreateInt(int val, LLVMContext &ctx) {
		IRBuilder<> builder(ctx);
		return builder.getInt(APInt(32,val));
	}

	// Create a function declaration.
	Function *ConstructTraceFunc(string name, LLVMContext &ctx, Module *module) {
		// Create a function with three strings and variable parameters like:
		// void lookin_enter(char *name, char *ty_desc, char *arg_desc, ...)
		vector<Type*> funcArgTy;
		funcArgTy.push_back(Type::getInt32Ty(ctx));
		funcArgTy.push_back(PointerType::get(Type::getInt8Ty(ctx),0));
		funcArgTy.push_back(PointerType::get(Type::getInt8Ty(ctx),0));
		funcArgTy.push_back(PointerType::get(Type::getInt8Ty(ctx),0));
		auto funTy = FunctionType::get(Type::getVoidTy(ctx), funcArgTy, true);
		auto funcVal = module->getOrInsertFunction(name, funTy).getCallee();
		if (funcVal == nullptr) {
			errs() << "Cannot construct trace func: " << name << "\n";
			return nullptr;
		}
		auto func = dyn_cast<Function>(funcVal);
		if (func == nullptr) {
			return nullptr;
		}
		return func;
	}

	string CppNamedDemangle(const char* abiName)
	{
		int status;
		char* name = abi::__cxa_demangle(abiName, nullptr, nullptr, &status);
		if (status != 0 || name == nullptr) {
			return abiName;
		}
		return name;
	}

	// Get a name of function argument.
	string FindArgName(Function &F, Argument &arg) {
		// Find dwarf data of function.
		auto findFunc = dwarfData.find(&F);
		if (findFunc == dwarfData.end())
			return "";

		// Find dwarf data of the argument.
		auto findArg = findFunc->second.find(&arg);
		if (findArg != findFunc->second.end())
			return findArg->second;

		// If not found, try to search def-use,
		// and try to find the argument stores to which local register.
		// The function is something like that:
		//
		// define dso_local i32 @gcd(i32 %0, i32 %1)
		// 	%3 = alloca i32, align 4
		// 	store i32 %0, i32* %3, align 4
		//      call void @llvm.dbg.declare(metadata i32* %3, metadata !15, metadata !DIExpression()), !dbg !16
		//
		// The argument %0 is stored to %3, then the name of %3 is the name of %0.
		if (!arg.hasOneUse())
			return "";
		auto use = arg.user_back();
		if (use == nullptr)
			return "";

		auto storeInst = dyn_cast<StoreInst>(use);
		if (storeInst == nullptr)
			return "";
		auto dest = storeInst->getPointerOperand();

		findArg = findFunc->second.find(dest);
		if (findArg == findFunc->second.end())
			return "";
		return findArg->second;
	}

	// Create a description string for function arguments.
	// For example, if function declaration is
	// 	int gcd(int a, int b)
	// then <typeDesc> is '22' as '2' means argument type is integer.
	// <nameDesc> is 'a b '.
	void CreateArgMeta(Function &F, Value **typeDesc, Value **nameDesc) {
		auto module = F.getParent();
                auto &ctx = F.getContext();
		string typeDescStr;
		string nameDescStr;
		for (auto &arg : F.args()) {
			auto ty = arg.getType();
			auto findTyChar = typeMap.find(ty->getTypeID());
			
			auto argName = FindArgName(F, arg);
			if (findTyChar == typeMap.end() && !argName.empty()) {
				typeDescStr += INVALID_CODE;
				nameDescStr += " ";
			} else {
				// Append type description and arg name.
				typeDescStr.append(1, findTyChar->second);
				nameDescStr += argName;
				nameDescStr += " ";
			}
		}
		*typeDesc = CreateString(typeDescStr, ctx, *module);
		*nameDesc = CreateString(nameDescStr, ctx, *module);
	}

	// Analyze DbgInfoIntrinsic and find name of a variable.
	void AnalyzeDwarfInst(Function &F, DbgInfoIntrinsic &dbgInst) {
		// <dbgInst> is something like:
		// 	call void @llvm.dbg.declare(metadata i32* %3, metadata !15, metadata !DIExpression()), !dbg !16
		//
		// The first operand is a Value*,
		// and the second one is a Meta data like:
		// 	!15 = !DILocalVariable(name: "a", arg: 1, scope: !11, file: !1, line: 5, type: !14)
		//
		// which is a DIVariable.
		//
		// Then we map %3 to the name 'a'.

		auto metaValue = cast<MetadataAsValue>(dbgInst.getOperand(0))->getMetadata();
		auto valueMeta = dyn_cast<ValueAsMetadata>(metaValue);
		if (valueMeta == nullptr)
			return;
		auto val = valueMeta->getValue();

		auto metaValue2 = dyn_cast<MetadataAsValue>(dbgInst.getOperand(1))->getMetadata();
		if (metaValue2 == nullptr)
			return;

		auto dwarf = dyn_cast<DIVariable>(metaValue2);
		if (dwarf == nullptr)
			return;
		dwarfData[&F][val] = dwarf->getName().str();
	}

	void AnalyzeDwarf(Function &F) {
		for (auto &bb : F) {
			for (auto &inst : bb) {
				auto dbgInst = dyn_cast<DbgInfoIntrinsic>(&inst);
				if (!dbgInst)
					continue;
				AnalyzeDwarfInst(F, *dbgInst);
			}
		}
	}

	// Insert a trace function to the beginning of function.
	bool InsertFuncEnter(Function &F, int funcIdx) {
		auto module = F.getParent();
		auto &ctx = F.getContext();
		IRBuilder<> builder(ctx);
		// Demangle the function name for easy reading.
		auto logName = CppNamedDemangle(F.getName().str().c_str());
		auto funcName = CreateString(logName, ctx, *module);
		if (funcName == nullptr) {
			errs() << "Cannot create func name variable: " << logName << "\n";
			return false;
		}
		Value *typeDesc = nullptr;
		Value *nameDesc = nullptr;
		CreateArgMeta(F, &typeDesc, &nameDesc);
		if (typeDesc == nullptr || nameDesc == nullptr) {
			errs() << "Cannot create type or name desc of func: " << logName << "\n";
			return false;
		}
		auto idxVal = CreateInt(funcIdx, ctx);

		// Prepare arguments.
                vector<Value *> argVec;
                argVec.push_back(idxVal);
                argVec.push_back(funcName);
                argVec.push_back(typeDesc);
                argVec.push_back(nameDesc);
		for (auto &arg : F.args()) {
                        argVec.push_back(&arg);
                }

		auto &bb = F.getEntryBlock();
                builder.SetInsertPoint(cast<Instruction>(F.getEntryBlock().begin()));
                builder.CreateCall(funcEnter, argVec);

		return true;
	}

	// Insert a trace function before RetInsts of the function.
	bool InsertFuncExit(Function &F, int funcIdx) {
		auto module = F.getParent();
                auto &ctx = F.getContext();
		IRBuilder<> builder(ctx);
		for (auto &bb : F) {
			for (auto &inst : bb) {
				auto retInst = dyn_cast<ReturnInst>(&inst);
				if (retInst == nullptr) {
					// Only care about ReturnInst
					continue;
				}

				// Create function name.
				auto logName = CppNamedDemangle(F.getName().str().c_str());
				auto msg = CreateString(logName, ctx, *module);
				if (msg == nullptr) {
					errs() << "Cannot create string value of func: " << logName << "\n";
					return false;
				}
				auto idxVal = CreateInt(funcIdx, ctx);

				vector<Value *> argVec;
				argVec.push_back(idxVal);
				argVec.push_back(msg);
				// Generate type and name desc for return value.
				auto retVal = retInst->getReturnValue();
				if (retVal && retVal->getType()->isIntegerTy()) {
					// TODO: care about other return types.
					auto typeDesc = CreateString(string(1, INT_CODE), ctx, *module);
					argVec.push_back(typeDesc);
					auto argDesc = CreateString("", ctx, *module);
					argVec.push_back(argDesc);
					argVec.push_back(retVal);
				} else {
					auto typeDesc = CreateString(string(1, INVALID_CODE), ctx, *module);
					argVec.push_back(typeDesc);
					auto argDesc = CreateString("", ctx, *module);
					argVec.push_back(argDesc);
				}

				builder.SetInsertPoint(cast<Instruction>(retInst));
				builder.CreateCall(funcExit, argVec);
			}
		}
		return true;
	}

	bool getFuncIdx(Function& F, int &idx) {
               int fd = open("/tmp/lookin_func_idx", O_RDWR);
                if (fd == -1) {
                        errs() << "Cannot open func idx file\n";
                        return false;
                }
                while (flock(fd, LOCK_EX) == -1);

                char buffer[16];
                auto bytes = read(fd, buffer, 16);
		if (bytes == -1) {
			close(fd);
			return false;
		}
		buffer[bytes] = '\0';
                idx = atoi(buffer);
		
                lseek(fd, 0, SEEK_SET);
                bytes = sprintf(buffer, "%d", idx+1);
		buffer[bytes] = '\0';
                write(fd, buffer, 16);

                while (flock(fd, LOCK_UN) == -1);
                close(fd);
                //errs() << "Func: " << F.getName() << " idx: " << idx << "\n";

		return true;
	}

	bool runOnFunction(Function& F) override {
		//    if (F.getName().str() != "Add")
		//		return false;
		//errs() << "Run on " << F.getName() << "\n";

		int idx = 0;
		if (!getFuncIdx(F, idx)) {
			errs() << "Cannot get function index\n";
			return false;
		}

		// Parse dwarf infos to get names of func arguments.
		AnalyzeDwarf(F);
		auto module = F.getParent();
		auto &ctx = F.getContext();
		if (traceFuncName.find(F.getName().str()) != traceFuncName.end()) {
			// Do not trace 'trace func' themselves.
			return false;
		}

		if (funcEnter == nullptr) {
			// Construct trace func declaration.
			funcEnter = ConstructTraceFunc("lookin_enter", ctx, module);
			funcExit = ConstructTraceFunc("lookin_exit", ctx, module);
		}

		if (!InsertFuncExit(F, idx)) {
			errs() << "Cannot insert trace exit func for: " << F.getName() << "\n";
		}
		if (!InsertFuncEnter(F, idx)) {
			errs() << "Cannot insert trace enter func for: " << F.getName() << "\n";
		}
		//errs().write_escaped(F.getName()) << "\n";
		//errs() << "Func: " << F << "\n";
		return false;
	}

	Function *funcEnter = nullptr;
	Function *funcExit = nullptr;
	set<string> traceFuncName = {"lookin_enter", "lookin_exit"};

	map<Function*, map<Value*, string>> dwarfData;

	const char FLT_CODE = '1';
	const char INT_CODE = '2';
	const char CHAR_CODE = '6';
	const char PTR_CODE = '7';
	const char INVALID_CODE = 'z';
	map<Type::TypeID, char> typeMap = {
		{Type::HalfTyID, FLT_CODE},
		{Type::BFloatTyID, FLT_CODE},
		{Type::FloatTyID, FLT_CODE},
		{Type::DoubleTyID, FLT_CODE},
		{Type::IntegerTyID, INT_CODE},
		{Type::PointerTyID, PTR_CODE},
	};

    }; 
}

char lookinPass::ID = 0;

static RegisterPass<lookinPass> X("lookin", "Lookin Pass",
		false /* Only looks at CFG */,
		false /* Analysis Pass */);
